package auth

import (
	"context"
	"fmt"

	"github.com/dgrijalva/jwt-go"
	"github.com/go-kratos/kratos/v2/errors"
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport/http"
)

// Claims 信息
type Claims struct {
	jwt.StandardClaims

	GID     string `json:"g_id"`
	BrandID int64  `json:"brand_id"`
	EmpID   int64  `json:"emp_id"`
	Name    string `json:"name"`
}
type contextKey string

const authKey = contextKey("auth")

const tokenHeader = "B-Token"

// FromContext 从上下文获取用户信息
func FromContext(ctx context.Context) (*Claims, error) {
	if claims, ok := ctx.Value(authKey).(*Claims); ok {
		return claims, nil
	}
	return nil, errors.Unauthorized("AUTH_ERROR", "认证错误")
}

// Admin is an admin auth middleware.
func Admin(key []byte) middleware.Middleware {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (interface{}, error) {
			if info, ok := http.FromServerContext(ctx); ok {
				r := info.Request
				bearToken := r.Header.Get(tokenHeader)

				if bearToken == "" {
					return nil, errors.Unauthorized("AUTH_ERROR", "认证错误")
				}
				token, err := new(jwt.Parser).ParseWithClaims(bearToken, &Claims{}, func(token *jwt.Token) (interface{}, error) {
					//Make sure that the token method conform to "SigningMethodHMAC"
					if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
						return nil, errors.Unauthorized("AUTH_ERROR", "认证错误")
					}
					key, err := jwt.ParseRSAPublicKeyFromPEM(key)
					if err != nil {
						return nil, fmt.Errorf("config jwt_pub_key invalid: %v", err)
					}

					return key, nil
				})
				if err != nil {
					return nil, errors.Unauthorized("AUTH_ERROR", "认证错误")
				}

				claims, ok := token.Claims.(*Claims)
				if !ok {
					return nil, errors.Unauthorized("AUTH_ERROR", "认证错误")
				}

				ctx = context.WithValue(ctx, authKey, claims)

			}
			reply, err := handler(ctx, req)
			if err != nil {
				return nil, err
			}

			return reply, nil
		}
	}
}
